import numpy as np
import os

# --- Main Configuration ---

# Define all parameters for the report
REPORT_PARAMETERS = {
    'definitions_designs': [
        (2, 11),
        (2, 12), 
        (2,  2),   
        (1, 11), 
        (1, 12), 
        (1,  2) # Each pair will generate a new \begin{table}
    ],
    'n_values': [100, 200, 400], # Sample sizes (table rows)
    'case_values': [1, 2],         # Case numbers (table panels)
    'input_folder': 'output' # Folder where .npz files are
}

# # Map case numbers to their LaTeX descriptions
# CASE_DESCRIPTIONS = {
#     1: r'$\epsilon\sim N(0,1)$ truncated at $[-3/2, 3/2]$',
#     2: r'$\epsilon\sim N(0,1.5)$ truncated at [-2,2]',
#     # 3: r'Another case description...'
# }

# The final .tex file that will be created
OUTPUT_LATEX_FILE = 'my_report.tex'

# --- Helper Functions (LaTeX String Generation) ---

def get_latex_preamble():
    """Returns the LaTeX document header."""
    return r"""\documentclass[11pt]{article}
\usepackage[utf8]{inputenc}
\usepackage[T1]{fontenc}
\usepackage{booktabs}
\usepackage{amsmath}
\usepackage[margin=1in]{geometry}
\usepackage{hyperref}

\begin{document}

"""

def get_latex_footer():
    """Returns the LaTeX document footer."""
    return r"""
\end{document}
"""

def get_table_start(defn, design):
    """Returns the LaTeX code for the start of a table."""
    return f"""
\\begin{{table}}[ht]
\\centering
Definition {defn}, Design {design}
\\bigskip
\\label{{tab:defn{defn}_design{design}}}
\\begin{{tabular}}{{l rrr rrr rrr}}
\\toprule
"""

def get_table_end():
    """Returns the LaTeX code for the end of a table."""
    return r"""
\bottomrule
\end{tabular}
\par \small
\end{table}
"""

def get_case_panel_header(case_num, case_desc):
    """Returns the header for a specific case panel."""
    return rf"""
% --- Case {case_num} Headers ---
& \multicolumn{{9}}{{c}}{{{case_desc}}} \\ [0.5em]
& \multicolumn{{3}}{{c}}{{$\hat{{\beta}}_{{x}}$}} & \multicolumn{{3}}{{c}}{{$\hat{{\beta}}_{{m}}$}} & \multicolumn{{3}}{{c}}{{$1-\dot{{\gamma}}$}} \\
\cmidrule(lr){{2-4}} \cmidrule(lr){{5-7}} \cmidrule(lr){{8-10}}
$n$ & \multicolumn{{1}}{{c}}{{Bias}} & \multicolumn{{1}}{{c}}{{Var}} & \multicolumn{{1}}{{c}}{{MSE}} & \multicolumn{{1}}{{c}}{{Bias}} & \multicolumn{{1}}{{c}}{{Var}} & \multicolumn{{1}}{{c}}{{MSE}} & \multicolumn{{1}}{{c}}{{Bias}} & \multicolumn{{1}}{{c}}{{Var}} & \multicolumn{{1}}{{c}}{{MSE}} \\
\midrule
"""

def format_data_row(n, bias, var, mse):
    """Formats the data for one row (one 'n') of the table."""
    # {bias[0]:.4f} formats the number to 4 decimal places
    # Estimator 1 (\beta_x): bias[0], var[0], mse[0]
    # Estimator 2 (\beta_m): bias[1], var[1], mse[1]
    # Estimator 3 (1-\gamma): bias[2], var[2], mse[2]
    return f"{n} & " \
           f"{bias[0]:.4f} & {var[0]:.4f} & {mse[0]:.4f} & " \
           f"{bias[1]:.4f} & {var[1]:.4f} & {mse[1]:.4f} & " \
           f"{bias[2]:.4f} & {var[2]:.4f} & {mse[2]:.4f} \\\\"

# --- Main Script Function ---

def generate_latex_report():
    """
    Main function to generate the complete LaTeX report.
    """
    print(f"Starting LaTeX report generation...")
    
    # Initialize the final LaTeX string with the preamble
    latex_content = get_latex_preamble()
    
    params = REPORT_PARAMETERS
    
    # --- 1. Loop over each (Definition, Design) pair ---
    # Each pair gets its own \begin{table}
    for defn, design in params['definitions_designs']:
        print(f"  Generating table for: Definition {defn}, Design {design}")
        
        latex_content += get_table_start(defn, design)

        if defn == 1:
            CASE_DESCRIPTIONS = {
                1: r'$\epsilon\sim N(0,1)$ truncated at $[-1.5, 1.5]$',
                2: r'$\epsilon\sim N(0,1.5)$ truncated at [-2,2]',
            }
        else: # defn == 2
            CASE_DESCRIPTIONS = {
                1: r'$\epsilon\sim N(0,1.5)$ truncated at $[-1.5, 1.5]$',
                2: r'$\epsilon\sim N(0,3)$ truncated at $[-3, 3]$',
            }
        
        # --- 2. Loop over each Case ---
        # Each case gets its own panel within the table
        for i, case in enumerate(params['case_values']):
            print(f"    - Adding panel for Case {case}")
            
            # Get the text description for this case
            case_desc = CASE_DESCRIPTIONS.get(case, f"Case {case}")
            latex_content += get_case_panel_header(case, case_desc)
            
            # --- 3. Loop over each n_value ---
            # Each n_value gets a row in the panel
            for n in params['n_values']:
                
                # Construct the expected filename
                filename = f'output_n{n}_defn{defn}_design{design}_case{case}.npz'
                filepath = os.path.join(params['input_folder'], filename)
                
                try:
                    # Load the .npz file
                    data = np.load(filepath)
                    bias = data['bias']
                    var = data['var']
                    mse = data['mse']
                    
                    # Add the formatted data row to the LaTeX string
                    latex_content += format_data_row(n, bias, var, mse) + "\n"
                
                except FileNotFoundError:
                    print(f"      [WARNING] File not found, skipping: {filepath}")
                    latex_content += f"% {n} & (File not found) & & & & & & & & \\\\\n"
                except Exception as e:
                    print(f"      [ERROR] Failed to load {filepath}: {e}")
                    latex_content += f"% {n} & (Error loading file) & & & & & & & & \\\\\n"

            # Add a \midrule between panels
            if i < len(params['case_values']) - 1:
                latex_content += "\\midrule\n"
        
        # Close the table environment
        latex_content += get_table_end()
        latex_content += "\n\n" # Add space between tables

    # Add the final document footer
    latex_content += get_latex_footer()
    
    # --- 4. Write the final string to the .tex file ---
    try:
        with open(OUTPUT_LATEX_FILE, 'w', encoding='utf-8') as f:
            f.write(latex_content)
        print(f"\nSuccess! LaTeX report saved to: {OUTPUT_LATEX_FILE}")
    except Exception as e:
        print(f"\n[ERROR] Failed to write output file: {e}")


# --- Run the script ---
if __name__ == "__main__":
    # First, check if the input folder exists
    if not os.path.isdir(REPORT_PARAMETERS['input_folder']):
        print(f"Error: Input folder '{REPORT_PARAMETERS['input_folder']}' not found.")
        print("Please run the 'create_dummy_data.py' script first, or update the 'input_folder' path.")
    else:
        generate_latex_report()